//
// Created by Tunm-Air13 on 2024/3/14.
//

#include <iostream>
#include "fasttext/fasttext.h"

int main() {
    // 创建一个fastText实例
    fasttext::FastText model;

    // 模型训练参数
    std::string inputFilePath = "../datasets/training_data.txt";
    std::string modelFileName = "td_dim32_b3_cpp";

    // 训练参数配置
    fasttext::Args args;
    args.input = inputFilePath;  // 训练数据文件路径
    args.output = modelFileName; // 不包括后缀的模型文件名
    args.epoch = 25;             // 训练轮数
    args.dim = 32;               // 向量维度
    args.lr = 1.0;               // 学习率
    args.model = fasttext::model_name::sup; // 设置为监督学习模式
    args.saveOutput = true;      // 保存输出
    args.output = modelFileName;

    // 训练模型
    model.train(args);

    // 保存模型
    model.saveModel(modelFileName + ".ftz");

    // 测试模型
    std::ifstream inFile("test_data.txt");
    int32_t k = 1; // 通常对于分类任务，k 设置为 1

    // 使用正确的 test 方法
    auto result = model.test(inFile, k);
    inFile.close();

    int64_t nexamples = std::get<0>(result);
    double precision = std::get<1>(result);
    double recall = std::get<2>(result);

    std::cout << "N: " << nexamples << ", Precision: " << precision << ", Recall: " << recall << std::endl;

    return 0;
}
